#ifndef SRC_POWEROF2_LINEAR_P2_CRYPTO_AGENT
#define SRC_POWEROF2_LINEAR_P2_CRYPTO_AGENT
#include "crypto_helper.hpp"

// * linearp2_hmmcs:
//   * linear : linear encoding
//   * p2     : m = power of 2
//   * hmm    : homomorphic matrix multiplication
//   * cs     : cryptograph server
class linearp2_hmmcs : public hmm_crypto_agent {
public:
  /**
   * @brief Construct a new `linearp2_hmmcs` object
   *
   * @tparam Scheme         helib::BGV or helib::CKKS
   * @param params          defualt parameters for `Scheme`
   * @param expected_nslot  expected # of slots
   */
  template <typename Scheme>
  linearp2_hmmcs(const params<Scheme> &params, long expected_nslot);

  /**
   * @brief Construct a new `linearp2_hmmcs` object
   *        Building the HElib context externally.
   *        * Must provide the building time
   * @param context HElib context
   * @param time    Escaping time of building HElib context
   */
  linearp2_hmmcs(std::shared_ptr<helib::Context> context, double time)
      : hmm_crypto_agent(context, time) {}

  // See below
  void report_time() const override;
  double setup_time() const { return setup_time_; }
  double find_params_time() const { return find_params_time_; }

protected:
  // * setup_time
  //   * find_params_time (m, plaintext modulus & generators)
  //   * build_context_time (HElib context) <- hmm_crypto_agent
  double setup_time_ = 0.0;
  double find_params_time_ = 0.0;
};

template <typename Scheme>
class linearp2_hmmcc : public linearp2_hmmcs, public hmm_crypto_client<Scheme> {
public:
  /**
   * @brief Construct a new `linearp2_hmmcc` object
   *        See the constructor of `linearp2_hmmcs` for more details
   */
  linearp2_hmmcc(const params<Scheme> &params, long expected_nslot);

  /**
   * @brief Construct a new `linearp2_hmmcc` object
   *        using a linearp2_hmmcs instance
   * @param hmmcs HMM cryptograph server
   */
  linearp2_hmmcc(const linearp2_hmmcc &hmmcs);

  // See below
  void report_time() const override;

  /**
   * @brief Implementation of the encryption method for `hmmcc`
   *        based on power-of-2 packing technology
   *        See `hmm_crypto_client` in crypto_helper.hpp for more details
   */
  helib::Ctxt encrypt(const NTL::mat_ZZ &mat, ENC type = PUBKEY) override;
  void encrypt(helib::Ctxt &ctxt, const NTL::mat_ZZ &mat, ENC type);
  helib::Ctxt encrypt(const std::vector<long> &vec, ENC type = PUBKEY) override;
  helib::Ctxt encrypt(const std::vector<double> &vec,
                      ENC type = PUBKEY) override;
  template <typename T>
  helib::Ctxt encrypt(T arr[], long len, ENC type = PUBKEY) {
    const long nslots = context_->getEA().size();
    ASSERT_PRINTF(nslots >= len,
                  "Array length %ld exceeds the number of slots %ld\n", len,
                  nslots);
    std::vector<T> vec(nslots);
    std::copy(arr, arr + len, vec.begin());
    return std::move(encrypt(vec, type));
  }

  /**
   * @brief Implementation of the decryption method for `hmmcc`
   *        based on power-of-2 packing technology
   *        See `hmm_crypto_client` in crypto_helper.hpp for more details
   */
  void decrypt(NTL::mat_ZZ &mat, const helib::Ctxt &cmat) override;
  void decrypt(std::vector<long> &vec, const helib::Ctxt &cmat) override;
  void decrypt(std::vector<double> &vec, const helib::Ctxt &cmat) override;
};

#endif /* SRC_POWEROF2_LINEAR_P2_CRYPTO_AGENT */
